import math
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
import numpy as np
import circular_inter_fast
reload(circular_inter_fast)


class LIP(object):
    # if the name of var is not world ,it means it is in local coodinate

    def __init__(self, x0, dx0, z, dt, stand_leg):
        self.x0 = x0
        self.dx0 = dx0
        self.x = x0
        self.dx = dx0
        self.z = z
        self.left_foot_x = 0.0
        self.left_foot_z = 0.0
        self.right_foot_x = 0.0
        self.right_foot_z = 0.0

        self.left_knee_x = 0.0
        self.left_knee_z = 0.0
        self.right_knee_x = 0.0
        self.right_knee_z = 0.0

        self.swing_foot_x = 0.0
        self.swing_foot_z = 0.0

        self.g = 9.8
        self.ddx = self.g / self.z * self.x
        self.Tc = math.sqrt(self.z / self.g)
        # t Will be set to 0,when enter a new state
        self.t = 0.0
        self.dt = dt
        self.stand_leg_world_x = 0.0
        self.E = 0.0
        self.capture_point = 0.0
        self.stand_leg = stand_leg
        self.left_leg_length = 0.0
        self.right_leg_length = 0.0
        self.f = 0.0
        self.inster_function = None
        self.shank_length = 0.4
        self.thigh_length = 0.4
        self.leg_length_max = self.thigh_length + self.shank_length
        self.init_figure()
        self.step_number = 0
        self.fall = 0
        self.out_range_left = 0
        self.out_range_right = 0
        self.out_range = 0
        self.m = 40

    def update(self):
        self.step_number = self.step_number + 1
        self.t = self.t + self.dt
        # self.update_motion_state()
        self.update_motion_state_recur(self.f)
        self.update_knee_pos()
        self.update_orbital_energy()
        self.update_leg_Length()

        if self.stand_leg == 'left_leg':
            x = self.right_foot_x
            z = self.right_foot_z
        else:
            x = self.left_foot_x
            z = self.left_foot_z
        observation = [self.x, self.dx, self.ddx, x, z]
        # print(observation)
        reward = 0.0
        done = 0

        self.out_range = self.out_range_left or self.out_range_right
        if self.step_number >= 5000:
            done = 1
        if self.fall == 1:
            reward = -10.0
            done = 2
        elif self.out_range == 1:
            done = 3
            reward = -5.0
        # else:
            # reward = self.step_number * 0.00001
        info = self.step_number
        # print('done', done, 'step_number', self.step_number,'fail',self.fall,'out_range',self.out_range)
        return np.array(observation), reward, done, info

    def update_orbital_energy(self):
        self.E = (self.dx**2 - (self.g / self.z) * self.x**2) / 2

    def update_motion_state(self):
        t = self.t
        tau = t / self.Tc
        self.ddx = self.g / self.z * self.x
        self.x = (self.x0 * math.cosh(tau) +
                   self.Tc * self.dx0 * math.sinh(tau))
        self.dx = ((self.x0 / self.Tc) * math.sinh(tau) +
                    self.dx0 * math.cosh(tau))
        self.ddx = self.g / self.z * self.x

    def update_motion_state_recur(self, f):
        self.ddx = self.g / self.z * self.x + self.f / self.m
        self.dx = self.ddx * self.dt + self.dx
        self.x = self.dx * self.dt + self.x

    def update_leg_Length(self):
        self.left_leg_length = math.sqrt((self.left_foot_x - self.x)**2 +
                                          (self.left_foot_z - self.z)**2)
        self.right_leg_length = math.sqrt((self.right_foot_x - self.x)**2 +
                                           (self.right_foot_z - self.z)**2)
        self.fall = 0
        if self.left_leg_length > self.leg_length_max and self.stand_leg == 'lef_leg':
            self.fall = 1
        if self.right_leg_length > self.leg_length_max and self.stand_leg == 'right_leg':
            self.fall = 1


    def switch_support_leg(self, stand_leg):
        # import ipdb; ipdb.set_trace()  # XXX BREAKPOINT
        if self.stand_leg != stand_leg:
            if self.stand_leg == 'left_leg':
                self.x0 = self.x - self.right_foot_x
                self.stand_leg = 'right_leg'
                self.stand_leg_world_x = self.stand_leg_world_x + self.right_foot_x
                # raw_input()
                # import ipdb; ipdb.set_trace()  # XXX BREAKPOINT

                self.left_foot_x = -self.right_foot_x
                self.right_foot_x = 0
            else:
                self.x0 = self.x - self.left_foot_x
                self.stand_leg = 'left_leg'
                self.stand_leg_world_x = self.stand_leg_world_x \
                    + self.left_foot_x
                self.right_foot_x = -self.left_foot_x
                self.left_foot_x = 0
            self.dx0 = self.dx
            self.x = self.x0

            self.t = 0

    def get_capture_point(self, E):
        cp = self.x + (
            self.dx / math.abs(self.dx)) * math.sqrt(
                (self.dx**2 - 2 * E) * self.z / self.g)
        return cp

    def set_swing_foot_pos(self, x, z):
        if self.stand_leg == 'left_leg':
            self.right_foot_x = x
            self.right_foot_z = z
        else:
            self.left_foot_x = x
            self.left_foot_z = z
        self.swing_foot_x = x
        self.swing_foot_z = z

    def set_external_force(self,f):
        self.f = f

    def reset(self, x0, dx0, z, dt, stand_leg):
        self.x0 = x0
        self.dx0 = dx0
        self.x = x0
        self.dx = dx0
        self.z = z
        self.left_foot_x = 0.0
        self.left_foot_z = 0.0
        self.right_foot_x = 0.0
        self.right_foot_z = 0.0

        self.left_knee_x = 0.0
        self.left_knee_z = 0.0
        self.right_knee_x = 0.0
        self.right_knee_z = 0.0

        self.g = 9.8
        self.ddx = self.g / self.z * self.x
        self.Tc = math.sqrt(self.z / self.g)
        # t Will be set to 0,when enter a new state
        self.t = 0.0
        self.dt = dt
        self.stand_leg_world_x = 0.0
        self.E = 0.0
        self.capture_point = 0.0
        self.stand_leg = stand_leg
        self.left_leg_length = 0.0
        self.right_leg_length = 0.0
        self.f = 0.0
        self.shank_length = 0.4
        self.thigh_length = 0.4
        self.leg_length_max = self.thigh_length + self.shank_length

        self.step_number = 0
        self.fall = 0
        self.out_range_left = 0
        self.out_range_right = 0
        self.out_range = 0

    def update_knee_pos(self):
        self.left_knee_x, self.left_knee_z, self.out_range_left \
            = circular_inter_fast.circular_inter_fast(self.x,
                                     self.z,
                                     self.thigh_length,
                                     self.left_foot_x,
                                     self.left_foot_z,
                                     self.shank_length,
                                     1)
        self.right_knee_x, self.right_knee_z, self.out_range_right = \
            circular_inter_fast.circular_inter_fast(self.x,
                                   self.z,
                                   self.thigh_length,
                                   self.right_foot_x,
                                   self.right_foot_z,
                                   self.shank_length,
                                   1)
    def update_plot(self, i):
        if self.inster_function!=None:
            self.inster_function(self)
        offset = self.stand_leg_world_x
        world_x = self.x + offset
        world_z = self.z
        world_left_foot_x = self.left_foot_x + offset
        world_left_foot_z = self.left_foot_z
        world_right_foot_x = self.right_foot_x + offset
        world_right_foot_z = self.right_foot_z

        world_left_knee_x = self.left_knee_x + offset
        world_left_knee_z = self.left_knee_z
        world_right_knee_x = self.right_knee_x + offset
        world_right_knee_z = self.right_knee_z

        self.body.center = world_x, world_z
        self.left_knee.center = world_left_knee_x, world_left_knee_z
        self.right_knee.center = (world_right_knee_x + 0.0), world_right_knee_z


        self.left_shank.set_data([world_left_foot_x, world_left_knee_x],
                                  [world_left_foot_z, world_left_knee_z])
        self.left_thigh.set_data([world_x, world_left_knee_x],
                                  [world_z, world_left_knee_z])
        self.right_shank.set_data([world_right_foot_x, world_right_knee_x],
                                   [world_right_foot_z, world_right_knee_z])
        self.right_thigh.set_data([world_x, world_right_knee_x],
                                   [world_z, world_right_knee_z])

        return self.body, self.left_knee, self.right_knee,\
            self.right_thigh,\
            self.left_thigh,\
            self.left_shank, self.right_shank

    def get_swing_foot_pos(self):
        if self.stand_leg == 'left_leg':
            x = self.right_foot_x
            z = self.right_foot_z
        else:
            x = self.left_foot_x
            z = self.left_foot_z
        return x, z

    def init_figure(self, blit=True):
        self.fig = plt.figure()
        self.ax = plt.axes(xlim=(-0.5, 3.5), ylim=(0, 1.5), aspect='equal')
        self.left_thigh = Line2D([], [], lw=2, marker='o',
                                  markersize=6, animated=blit, axes=self.ax)
        self.left_shank = Line2D([], [], lw=2, marker='o',
                                  markersize=6, animated=blit, axes=self.ax)
        self.right_thigh = Line2D([], [], lw=2, marker='o',
                                   markersize=6, animated=blit, axes=self.ax)
        self.right_shank = Line2D([], [], lw=2, marker='o',
                                   markersize=6, animated=blit, axes=self.ax)

        self.body = Circle((0, 0), 0.1, color='red', axes=self.ax,
                            animated=blit)
        self.left_knee = Circle((0.0, 0), 0.07, color='red', axes=self.ax,
                                 animated=blit)
        self.right_knee = Circle((0.0, 0), 0.07, color='red', axes=self.ax,
                                  animated=blit)

        self.ax.add_artist(self.right_thigh)
        self.ax.add_artist(self.right_shank)
        self.ax.add_artist(self.left_thigh)
        self.ax.add_artist(self.left_shank)
        self.ax.add_patch(self.left_knee)
        self.ax.add_patch(self.right_knee)
        self.ax.add_patch(self.body)

    def run(self, blit=True):

        # anim = animation.FuncAnimation(fig, self.updateplot,
        # interval=20, blit=False, repeat=False)
        anim = animation.FuncAnimation(self.fig, self.update_plot,
                                       interval=1, blit=blit, repeat=False)
        plt.show()
        return anim

if __name__ == "__main__":
    lip = LIP(0.1, 0.1, 0.7, 0.001)
    lip.set_swing_foot_pos(0.1, 0.1)
    lip.run()
